#include <opencv2/opencv.hpp>
#include "calibration.hpp"
#include "config.hpp"

int main(int argc, char **argv)
{
    Size boardSize, imageSize;
    float squareSize, aspectRatio = 1;
    Mat cameraMatrix, distCoeffs;
    string outputFilename;
    string inputFilename = "";

    int i, nframes;
    bool writeExtrinsics, writePoints;
    bool undistortImage = false;
    int flags = 0;
    unique_ptr<VideoCapture> capture(new VideoCapture(1));

    bool flipVertical;
    bool showUndistorted;
    bool videofile;
    int delay;
    clock_t prevTimestamp = 0;
    int mode = DETECTION;
    int cameraId = 0;
    vector<vector<Point2f>> imagePoints;
    vector<string> imageList;
    Pattern pattern = CHESSBOARD;

    cv::CommandLineParser parser(argc, argv,
                                 "{help ||}{w||}{h||}{pt|chessboard|}{n|10|}{d|1000|}{s|1|}{o|out_camera_data.yml|}"
                                 "{op||}{oe||}{zt||}{a||}{p||}{v||}{V||}{su||}"
                                 "{oo||}{ws|11|}{dt||}"
                                 "{@input_data|0|}");
    if (parser.has("help"))
    {
        help();
        return 0;
    }
    boardSize.width = parser.get<int>("w");
    boardSize.height = parser.get<int>("h");
    if (parser.has("pt"))
    {
        string val = parser.get<string>("pt");
        if (val == "circles")
            pattern = CIRCLES_GRID;
        else if (val == "acircles")
            pattern = ASYMMETRIC_CIRCLES_GRID;
        else if (val == "chessboard")
            pattern = CHESSBOARD;
        else
            return fprintf(stderr, "Invalid pattern type: must be chessboard or circles\n"), -1;
    }
    squareSize = parser.get<float>("s");
    nframes = parser.get<int>("n");
    delay = parser.get<int>("d");
    writePoints = parser.has("op");
    writeExtrinsics = parser.has("oe");
    bool writeGrid = parser.has("oo");
    if (parser.has("a"))
    {
        flags |= CALIB_FIX_ASPECT_RATIO;
        aspectRatio = parser.get<float>("a");
    }
    if (parser.has("zt"))
        flags |= CALIB_ZERO_TANGENT_DIST;
    if (parser.has("p"))
        flags |= CALIB_FIX_PRINCIPAL_POINT;
    flipVertical = parser.has("v");
    videofile = parser.has("V");
    if (parser.has("o"))
        outputFilename = parser.get<string>("o");
    showUndistorted = parser.has("su");
    if (isdigit(parser.get<string>("@input_data")[0]))
        cameraId = parser.get<int>("@input_data");
    else
        inputFilename = parser.get<string>("@input_data");
    int winSize = parser.get<int>("ws");
    float grid_width = squareSize * (boardSize.width - 1);
    bool release_object = false;
    if (parser.has("dt"))
    {
        grid_width = parser.get<float>("dt");
        release_object = true;
    }
    if (!parser.check())
    {
        help();
        parser.printErrors();
        return -1;
    }
    if (squareSize <= 0)
        return fprintf(stderr, "Invalid board square width\n"), -1;
    if (nframes <= 3)
        return printf("Invalid number of images\n"), -1;
    if (aspectRatio <= 0)
        return printf("Invalid aspect ratio\n"), -1;
    if (delay <= 0)
        return printf("Invalid delay\n"), -1;
    if (boardSize.width <= 0)
        return fprintf(stderr, "Invalid board width\n"), -1;
    if (boardSize.height <= 0)
        return fprintf(stderr, "Invalid board height\n"), -1;

    if (!inputFilename.empty())
    {
        if (!videofile && readStringList(samples::findFile(inputFilename), imageList))
            mode = CAPTURING;
        // capture= RMVideoCapture();
    }
    // else
    // capture = RMVideoCapture();

    if (!capture->isOpened() && imageList.empty())
        return fprintf(stderr, "Could not initialize video (%d) capture\n", cameraId), -2;

    if (!imageList.empty())
        nframes = (int)imageList.size();

    if (capture->isOpened())
        printf("%s", liveCaptureHelp);

    namedWindow("Image View", 1);

    for (i = 0;; i++)
    {
        Mat view, viewGray;
        bool blink = false;

        if (capture->isOpened())
        {
            Mat view0;
            capture->read(view0);
            view0.copyTo(view);
        }
        else if (i < (int)imageList.size())
            view = imread(imageList[i], 1);

        if (view.empty())
        {
            cout << "dshfaldf" << endl;
            if (imagePoints.size() > 0)
                runAndSave(outputFilename, imagePoints, imageSize,
                           boardSize, pattern, squareSize, grid_width, release_object, aspectRatio,
                           flags, cameraMatrix, distCoeffs,
                           writeExtrinsics, writePoints, writeGrid);
            continue;
        }

        imageSize = view.size();

        if (flipVertical)
            flip(view, view, 0);

        vector<Point2f> pointbuf;
        cvtColor(view, viewGray, COLOR_BGR2GRAY);

        bool found;
        switch (pattern)
        {
        case CHESSBOARD:
            found = findChessboardCorners(view, boardSize, pointbuf,
                                          CALIB_CB_ADAPTIVE_THRESH | CALIB_CB_FAST_CHECK | CALIB_CB_NORMALIZE_IMAGE);
            break;
        case CIRCLES_GRID:
            found = findCirclesGrid(view, boardSize, pointbuf);
            break;
        case ASYMMETRIC_CIRCLES_GRID:
            found = findCirclesGrid(view, boardSize, pointbuf, CALIB_CB_ASYMMETRIC_GRID);
            break;
        default:
            return fprintf(stderr, "Unknown pattern type\n"), -1;
        }

        // improve the found corners' coordinate accuracy
        if (pattern == CHESSBOARD && found)
            cornerSubPix(viewGray, pointbuf, Size(winSize, winSize),
                         Size(-1, -1), TermCriteria(TermCriteria::EPS + TermCriteria::COUNT, 30, 0.0001));

        if (mode == CAPTURING && found &&
            (!capture->isOpened() || clock() - prevTimestamp > delay * 1e-3 * CLOCKS_PER_SEC))
        {
            imagePoints.push_back(pointbuf);
            prevTimestamp = clock();
            blink = capture->isOpened();
        }

        if (found)
            drawChessboardCorners(view, boardSize, Mat(pointbuf), found);

        string msg = mode == CAPTURING ? "100/100" : mode == CALIBRATED ? "Calibrated"
                                                                        : "Press 'g' to start";
        int baseLine = 0;
        Size textSize = getTextSize(msg, 1, 1, 1, &baseLine);
        Point textOrigin(view.cols - 2 * textSize.width - 10, view.rows - 2 * baseLine - 10);

        if (mode == CAPTURING)
        {
            if (undistortImage)
                msg = format("%d/%d Undist", (int)imagePoints.size(), nframes);
            else
                msg = format("%d/%d", (int)imagePoints.size(), nframes);
        }

        putText(view, msg, textOrigin, 1, 1,
                mode != CALIBRATED ? Scalar(0, 0, 255) : Scalar(0, 255, 0));

        if (blink)
            bitwise_not(view, view);

        if (mode == CALIBRATED && undistortImage)
        {
            Mat temp = view.clone();
            undistort(temp, view, cameraMatrix, distCoeffs);
        }

        imshow("Image View", view);
        char key = (char)waitKey(capture->isOpened() ? 50 : 500);

        if (key == 27)
            break;

        if (key == 'u' && mode == CALIBRATED)
            undistortImage = !undistortImage;

        if (capture->isOpened() && key == 'g')
        {
            mode = CAPTURING;
            imagePoints.clear();
        }

        if (mode == CAPTURING && imagePoints.size() >= (unsigned)nframes)
        {
            if (runAndSave(outputFilename, imagePoints, imageSize,
                           boardSize, pattern, squareSize, grid_width, release_object, aspectRatio,
                           flags, cameraMatrix, distCoeffs,
                           writeExtrinsics, writePoints, writeGrid))
                mode = CALIBRATED;
            else
                mode = DETECTION;
            if (!capture->isOpened())
                break;
        }
    }

    if (!capture->isOpened() && showUndistorted)
    {
        Mat view, rview, map1, map2;
        initUndistortRectifyMap(cameraMatrix, distCoeffs, Mat(),
                                getOptimalNewCameraMatrix(cameraMatrix, distCoeffs, imageSize, 1, imageSize, 0),
                                imageSize, CV_16SC2, map1, map2);

        for (i = 0; i < (int)imageList.size(); i++)
        {
            view = imread(imageList[i], 1);
            if (view.empty())
                continue;
            //undistort( view, rview, cameraMatrix, distCoeffs, cameraMatrix );
            remap(view, rview, map1, map2, INTER_LINEAR);
            imshow("Image View", rview);
            char c = (char)waitKey();
            if (c == 27 || c == 'q' || c == 'Q')
                break;
        }
    }

    return 0;
}